import tensorflow as tf
import numpy as np
from utils.config_utils import parse_anchors

YOLOv3_anchors = parse_anchors('config/anchor')


def bbox_to_target(bbox, label, input_shape=(416, 416), anchors=YOLOv3_anchors, num_classes=80):
    # bbox.shape = (box num, 4) 真实值坐标表示为(xmin,ymin,xmax,ymax)，是相对坐标
    # label.shape = (box num) 每个标注框的类别
    # anchors = (9,2)
    # 获取有几个尺度的输出，每个尺度对应3个anchor:3
    num_layers = anchors.shape[0] // 3
    # anchor对应的特征图掩码：第一个特征图对应第6，7，8个anchor...
    anchor_mask = tf.cond(tf.equal(num_layers, 3), lambda: tf.constant(
        [[6, 7, 8], [3, 4, 5], [0, 1, 2]]), lambda: tf.constant([[3, 4, 5], [1, 2, 3]]))

    # bbox的相对中心点坐标
    true_boxes_xy = (bbox[..., 0:2] + bbox[..., 2:4]) / 2.
    # bbox的相对宽高
    true_boxes_wh = tf.math.abs(bbox[..., 2:4] - bbox[..., 0:2])
    # bbox的结果:将中心点坐标和宽高拼接在一起
    true_boxes = tf.concat([true_boxes_xy, true_boxes_wh], axis=-1)
    # bbox的绝对坐标和绝对宽高
    #     boxes_xy = true_boxes[..., 0:2] * input_shape
    #     boxes_wh = true_boxes[..., 2:4] * input_shape  # box absolute size

    boxes_xy = true_boxes[..., 0:2]
    boxes_wh = true_boxes[..., 2:4]
    # 生成与yoloV3输出结果相同大小的全0数组：y_true.shape[layer] = (height, width, anchor num, 5 + class num)
    y_true = tuple((np.zeros(shape=(input_shape[0] // {0: 32, 1: 16, 2: 8}[l], input_shape[1] // {0: 32, 1: 16, 2: 8}[
        l], tf.shape(anchor_mask[l, ...])[0], 5 + num_classes), dtype=np.float32) for l in range(num_layers)))

    # 扩展一个维度，用来存放anchor的索引
    anchors = tf.expand_dims(
        tf.convert_to_tensor(anchors, dtype=tf.float32), 0)
    # 用于计算交并比
    # 以anchor中心为原点，计算右下角坐标
    anchor_maxes = anchors / 2.
    # 以anchor中心为原点，计算左上角坐标
    anchor_mins = -anchor_maxes
    # 创建一个mask,指明目标是否存在，宽度大于0的认为是真实的目标
    valid_mask = tf.greater(boxes_wh[..., 0], 0)
    #  wh.shape = (valid box num, 2)
    # 获取真实的目标的宽高
    wh = tf.boolean_mask(boxes_wh, valid_mask)
    # 获取真实目标的box：valid_true_boxes.shape = (valid box num, 4)
    valid_true_boxes = tf.boolean_mask(true_boxes, valid_mask)
    # 获取真实目标的标签值：valid_label.shape = (valid box num)
    valid_label = tf.boolean_mask(label, valid_mask)
    # if there is any valid bbox, get anchor box which has the maximum iou with current bbox.
    # 当图像中存在目标时，计算与目标交并比最大的anchor作为正样本，并设置标记结果
    if wh.shape[0] > 0:
        # 扩展一个维度，用来存放对应的anchor：wh.shape = (valid box num, 1, 2)
        wh = tf.expand_dims(wh, -2)
        # 以box的中心点为原点：计算右下角坐标：max of width, height, box_maxes.shape = (valid box num, 1, 2)
        box_maxes = wh / 2
        # 以box的中心点为原点：计算左上角坐标：min of width, height, box_mins.shape = (valid box num, 1, 2)
        box_mins = -box_maxes
        # 计算box与anchor交的左上角坐标：intersect_mins.shape = (valid box num, anchor num(9), 2)
        intersect_mins = tf.math.maximum(box_mins, anchor_mins)
        # 计算box与anchor交的右下角坐标：intersect_maxes.shape = (valid box num, anchor num(9), 2)
        intersect_maxes = tf.math.minimum(box_maxes, anchor_maxes)
        # 计算交集的宽高：intersect_wh.shape = (valid box num, anchor num(9), 2)
        intersect_wh = tf.math.maximum(intersect_maxes - intersect_mins, 0.)
        # 计算交集的面积：intersect_area.shape = (valid box num, anchor num(9))
        intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
        # 计算box的面积：box_area.shape = (valid box_num, 1)
        box_area = wh[..., 0] * wh[..., 1]
        # 计算anchor的面积：anchor_area.shape = (1, anchor num(9))
        anchor_area = anchors[..., 0] * anchors[..., 1]
        # 计算交并比：iou.shape = (valid box num, anchor num(9))
        iou = intersect_area / (box_area + anchor_area - intersect_area)
        # 计算与box交并比最大的anchor,将其作为正样本：best_anchor.shape = (valid box num)
        best_anchor = tf.math.argmax(iou, axis=-1, output_type=tf.int32)
        # 遍历与box匹配成功的anchor
        for t in range(tf.shape(best_anchor)[0]):
            # 获取第t个anchor
            n = best_anchor[t]
            # 获取anchor的位置
            pos = tf.where(tf.equal(anchor_mask, n))
            # 获取尺度值：0，1，2
            l = pos[0][0]
            # 获取对应的anchor索引
            k = pos[0][1]
            # 获取anchor对应的grid cell的行数，限制在0到最大值之间
            i = int(tf.clip_by_value(
                valid_true_boxes[t, 1] * y_true[l].shape[0], clip_value_min=0, clip_value_max=y_true[l].shape[0] - 1))
            # 获取anchor对应的grid cell的列数，限制在0到最大值之间
            j = int(tf.clip_by_value(
                valid_true_boxes[t, 0] * y_true[l].shape[1], clip_value_min=0, clip_value_max=y_true[l].shape[1] - 1))
            # 获取anchor的类别
            c = valid_label[t]
            # box的位置:(x,y,width,height)
            y_true[l][i, j, k, 0:4] = valid_true_boxes[t, 0:4]
            # 匹配上的都包含目标，置信度设为1
            y_true[l][i, j, k, 4] = 1
            # 类别信息
            y_true[l][i, j, k, 5 + c] = 1
    # 返回3个尺度对应的真实值
    return (tf.convert_to_tensor(y_true[0]), tf.convert_to_tensor(y_true[1]), tf.convert_to_tensor(y_true[2]))



if __name__ == "__main__":
    from dataset.get_tfdata import getdata

    datasets = getdata("../dataset/voc_val.tfrecords")

    for image, width, height, boxes, labels in datasets.take(1):
        # plt.imshow(image)
        label1, label2, label3 = bbox_to_target(bbox=boxes, label=labels, num_classes=20)
